import os
import shutil
from pathlib import Path


def rename_images_to_sequential(folder_path):
    """
    将文件夹中的图片重命名为1.jpg, 2.jpg, 3.jpg...的格式

    Args:
        folder_path (str): 图片文件夹路径
    """
    # 支持的图片格式
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}

    # 转换为Path对象
    folder = Path(folder_path)

    if not folder.exists():
        print(f"错误: 文件夹 {folder_path} 不存在!")
        return

    # 获取所有图片文件
    image_files = []
    for file in folder.iterdir():
        if file.is_file() and file.suffix.lower() in image_extensions:
            image_files.append(file)

    if not image_files:
        print(f"文件夹 {folder_path} 中没有找到图片文件!")
        return

    # 按文件名排序（保持一定的顺序）
    image_files.sort(key=lambda x: x.name.lower())

    print(f"找到 {len(image_files)} 张图片")
    print("开始重命名...")

    # 创建临时文件夹避免命名冲突
    temp_folder = folder / "temp_rename"
    temp_folder.mkdir(exist_ok=True)

    try:
        # 第一步：将所有文件移动到临时文件夹并重命名
        for i, old_file in enumerate(image_files, 1):
            # 保持原始扩展名，但统一为.jpg（如果需要的话）
            # 如果要保持原扩展名，使用: new_name = f"{i}{old_file.suffix}"
            new_name = f"{i}.jpg"
            temp_file = temp_folder / new_name

            # 移动并重命名到临时文件夹
            shutil.move(str(old_file), str(temp_file))
            print(f"  {old_file.name} -> {new_name}")

        # 第二步：将文件从临时文件夹移回原文件夹
        for temp_file in temp_folder.iterdir():
            if temp_file.is_file():
                final_file = folder / temp_file.name
                shutil.move(str(temp_file), str(final_file))

        # 删除临时文件夹
        temp_folder.rmdir()

        print(f"\n✅ 成功重命名 {len(image_files)} 张图片!")
        print(f"图片已重命名为: 1.jpg, 2.jpg, 3.jpg, ..., {len(image_files)}.jpg")

    except Exception as e:
        print(f"❌ 重命名过程中出现错误: {e}")
        # 尝试清理临时文件夹
        if temp_folder.exists():
            try:
                # 将临时文件夹中的文件移回原文件夹
                for temp_file in temp_folder.iterdir():
                    if temp_file.is_file():
                        shutil.move(str(temp_file), str(folder / temp_file.name))
                temp_folder.rmdir()
            except:
                print("⚠️  请手动清理temp_rename文件夹")


def rename_images_keep_extension(folder_path):
    """
    将文件夹中的图片重命名为1.原扩展名, 2.原扩展名...的格式
    保持原始文件扩展名
    """
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp'}

    folder = Path(folder_path)

    if not folder.exists():
        print(f"错误: 文件夹 {folder_path} 不存在!")
        return

    # 获取所有图片文件
    image_files = []
    for file in folder.iterdir():
        if file.is_file() and file.suffix.lower() in image_extensions:
            image_files.append(file)

    if not image_files:
        print(f"文件夹 {folder_path} 中没有找到图片文件!")
        return

    # 按文件名排序
    image_files.sort(key=lambda x: x.name.lower())

    print(f"找到 {len(image_files)} 张图片")
    print("开始重命名（保持原扩展名）...")

    # 创建临时文件夹
    temp_folder = folder / "temp_rename"
    temp_folder.mkdir(exist_ok=True)

    try:
        # 重命名并移动到临时文件夹
        for i, old_file in enumerate(image_files, 1):
            new_name = f"{i}{old_file.suffix}"  # 保持原扩展名
            temp_file = temp_folder / new_name

            shutil.move(str(old_file), str(temp_file))
            print(f"  {old_file.name} -> {new_name}")

        # 移回原文件夹
        for temp_file in temp_folder.iterdir():
            if temp_file.is_file():
                final_file = folder / temp_file.name
                shutil.move(str(temp_file), str(final_file))

        temp_folder.rmdir()

        print(f"\n✅ 成功重命名 {len(image_files)} 张图片!")

    except Exception as e:
        print(f"❌ 重命名过程中出现错误: {e}")


if __name__ == "__main__":
    # 你的文件夹路径
    folder_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\测试\images"

    print("选择重命名方式:")
    print("1. 统一重命名为 1.jpg, 2.jpg, 3.jpg...")
    print("2. 保持原扩展名 1.png, 2.jpg, 3.bmp...")

    choice = input("请输入选择 (1 或 2): ").strip()

    if choice == "1":
        rename_images_to_sequential(folder_path)
    elif choice == "2":
        rename_images_keep_extension(folder_path)
    else:
        print("无效选择，默认使用方式1")
        rename_images_to_sequential(folder_path)

    input("\n按回车键退出...")